import pandas as pd
import numpy as np
from datetime import datetime
import matplotlib
matplotlib.use('Agg')  # 设置后端
import matplotlib.pyplot as plt
import mplcyberpunk
from matplotlib import font_manager
from .sentinel import ArbitrageSentinel
from .config import (ANALYSIS_CONFIG, FONT_PATH, 
                    FUTURES_CONFIG)  # 添加 FUTURES_CONFIG 导入

class ArbitrageBacktest:
    def __init__(self, initial_capital=1000000):
        """
        初始化回测系统
        :param initial_capital: 初始资金
        """
        self.initial_capital = initial_capital
        self.reset()
        
        # 风险控制参数
        self.max_position_ratio = 0.3  # 单个品种最大持仓比例
        self.max_total_position_ratio = 0.6  # 总持仓上限
        self.stop_loss_ratio = 0.02  # 单笔止损比例
        self.max_drawdown_limit = 0.1  # 最大回撤限制
        
        # 使用配置文件中的期货参数
        self.futures_config = FUTURES_CONFIG
        
        # 设置matplotlib样式
        plt.style.use("cyberpunk")
        
        # 设置中文字体
        self.font_prop = font_manager.FontProperties(fname=FONT_PATH)
        plt.rcParams['font.family'] = self.font_prop.get_name()
        
        # 初始化哨兵系统
        self.sentinel = ArbitrageSentinel(
            std_threshold=2.0,
            window_sizes=[
                ANALYSIS_CONFIG['windows']['short'],
                ANALYSIS_CONFIG['windows']['medium'],
                ANALYSIS_CONFIG['windows']['long']
            ]
        )
        
        # 交易参数
        self.position_size = 0.1  # 每次交易使用资金的比例
        self.max_positions = 5    # 最大同时持仓数
        
    def reset(self):
        """重置回测状态"""
        self.capital = self.initial_capital
        self.available_capital = self.initial_capital
        self.positions = {}
        self.trades = []
        self.daily_capital = []
        
    def run_backtest(self, price_data, pair_name):
        """运行回测"""
        self.reset()
        
        results = []
        equity_curve = []
        current_position = None
        last_force_close_direction = None  # 记录上次强制平仓的方向
        
        for i in range(max(self.sentinel.window_sizes), len(price_data)):
            date = price_data.index[i]
            current_price = price_data['收盘'].iloc[i]
            
            # 使用前i天的数据进行分析
            historical_data = price_data.iloc[:i+1]
            signal = self.sentinel.analyze_spread(historical_data, pair_name)
            
            # 如果有持仓，检查是否需要调整仓位或平仓
            if current_position is not None:
                # 记录持仓状态
                position_record = self._record_position_status(current_position, current_price, date)
                results.append(position_record)
                
                # 检查是否需要调整仓位或平仓
                exit_signal = self._check_exit_signals(current_position, current_price)
                if exit_signal[0]:
                    if "加仓" in exit_signal[1] or "平掉" in exit_signal[1]:
                        current_position, trade_record = self._handle_position_adjustment(
                            current_position, current_price, exit_signal[1], date
                        )
                        if trade_record:
                            results.append(trade_record)
                    else:  # 全部平仓
                        trade_result = self._close_position(
                            current_position, current_price, date, exit_signal[1]
                        )
                        results.append(trade_result)
                        
                        # 如果是强制平仓，记录方向
                        if "全局风险控制止损" in exit_signal[1]:
                            last_force_close_direction = current_position['direction']
                        
                        current_position = None
            
            # 如果没有持仓且有信号，检查是否可以开新仓
            elif signal['stop_settings'] is not None:
                new_direction = signal['stop_settings']['direction']
                
                # 检查是否违反强制平仓后的方向限制
                if last_force_close_direction is not None:
                    if new_direction == last_force_close_direction:
                        continue  # 跳过同向开仓
                    else:
                        last_force_close_direction = None  # 重置方向限制
                
                current_position = self.open_position(
                    price_data=historical_data,
                    entry_date=date,
                    direction=new_direction,
                    pair_name=pair_name
                )
                if current_position:
                    results.append(self._record_open_position(current_position))
            
            # 记录每日权益
            equity = self.calculate_current_equity(current_position, current_price)
            equity_curve.append({'date': date, 'equity': equity})
        
        return self._generate_backtest_report(results, equity_curve, pair_name)
    
    def _check_risk_limits(self, pair_name, position_size):
        """检查风险限制"""
        # 检查单个品种持仓限制
        if position_size > self.capital * self.max_position_ratio:
            return False, "超过单个品种持仓限制"
        
        # 检查总持仓限制
        total_position = sum(pos['margin'] for pos in self.positions.values())
        if (total_position + position_size) > self.capital * self.max_total_position_ratio:
            return False, "超过总持仓限制"
        
        # 检查可用资金
        if position_size > self.available_capital:
            return False, "可用资金不足"
        
        return True, ""

    def calculate_position_size(self, capital, price, pair_name):
        """
        计算开仓手数
        :param capital: 当前权益
        :param price: 开仓价格
        :param pair_name: 期货对名称
        :return: 开仓手数
        """
        try:
            # 如果资金为负，不允许开仓
            if capital <= 0:
                return 0
            
            # 获取期货配置
            future_config = FUTURES_CONFIG.get(pair_name)
            if not future_config:
                return 0
            
            # 获取合约参数
            margin_ratio = future_config['margin_ratio']  # 保证金比例
            point_value = future_config['point_value']    # 每点价值
            
            # 计算单手合约价值
            contract_value = price * point_value
            
            # 计算单手所需保证金
            margin_per_lot = contract_value * margin_ratio
            
            # 使用当前权益的10%计算可开手数
            position_capital = capital * 0.1  # 使用10%的权益
            max_lots = int(position_capital / margin_per_lot)
            
            # 确保不超过最大限制
            max_lots = min(max_lots, future_config['max_order'])
            
            # 确保至少有1手，如果资金不够1手则返回0
            return max_lots if max_lots >= 1 else 0
            
        except Exception as e:
            print(f"计算仓位出错: {str(e)}")
            return 0

    def open_position(self, price_data, entry_date, direction, pair_name):
        """开仓"""
        # 如果权益为负，不允许开仓
        if self.capital <= 0:
            return None
        
        entry_price = price_data.loc[entry_date, '收盘']
        
        # 计算开仓手数
        size = self.calculate_position_size(
            capital=self.available_capital,
            price=entry_price,
            pair_name=pair_name
        )
        
        # 如果计算出的手数为0，不开仓
        if size == 0:
            return None
        
        # 获取期货配置
        future_config = FUTURES_CONFIG.get(pair_name)
        margin_ratio = future_config['margin_ratio']
        point_value = future_config['point_value']
        
        # 计算保证金
        margin = entry_price * point_value * size * margin_ratio
        
        # 检查是否有足够的可用资金
        if margin > self.available_capital:
            return None
        
        # 使用哨兵系统计算止损止盈价位
        signal = self.sentinel.analyze_spread(price_data, pair_name)
        if signal['stop_settings'] is None:
            return None
        
        # 获取网格大小
        grid_size = signal['stop_settings']['grid_size']
        
        # 更新可用资金
        self.available_capital -= margin
        
        # 记录持仓
        position = {
            'entry_date': entry_date,
            'entry_price': entry_price,
            'direction': direction,
            'lots': size,
            'margin': margin,
            'pair_name': pair_name,
            'grid_size': grid_size,  # 添加网格大小
            'first_add': False,      # 添加加仓状态标记
            'second_add': False,
            'first_add_price': None,
            'first_add_lots': 0,
            'second_add_price': None,
            'second_add_lots': 0
        }
        
        self.positions[pair_name] = position
        return position

    def _check_exit_signals(self, position, current_price):
        """检查是否触发止损止盈"""
        # 计算当前权益
        current_equity = self.calculate_current_equity(position, current_price)
        
        # 检查全局风险控制 - 权益下降30%
        if current_equity <= self.initial_capital * 0.7:
            return current_price, "全局风险控制止损"
        
        # 获取网格大小
        grid_size = position['grid_size']
        
        # 计算止盈价位
        if position['direction'] == '多头':
            take_profit_1 = position['entry_price'] + grid_size * 3  # 3个网格止盈1
            take_profit_2 = position['entry_price'] + grid_size * 5  # 5个网格止盈2
            
            # 检查止盈
            if current_price >= take_profit_2:
                return current_price, "止盈2"
            elif current_price >= take_profit_1:
                return current_price, "止盈1"
            
            # 检查是否需要加仓（多头）
            if position.get('second_add', False):
                # 如果价格回到第一次加仓价格，平掉第二次加仓的仓位
                if current_price >= position['first_add_price']:
                    return current_price, "平掉第二次加仓"
            elif position.get('first_add', False):
                # 如果价格回到开仓价格，平掉第一次加仓的仓位
                if current_price >= position['entry_price']:
                    return current_price, "平掉第一次加仓"
        
        else:  # 空头
            take_profit_1 = position['entry_price'] - grid_size * 3  # 3个网格止盈1
            take_profit_2 = position['entry_price'] - grid_size * 5  # 5个网格止盈2
            
            # 检查止盈
            if current_price <= take_profit_2:
                return current_price, "止盈2"
            elif current_price <= take_profit_1:
                return current_price, "止盈1"
            
            # 检查是否需要加仓（空头）
            if position.get('second_add', False):
                # 如果价格回到第一次加仓价格，平掉第二次加仓的仓位
                if current_price <= position['first_add_price']:
                    return current_price, "平掉第二次加仓"
            elif position.get('first_add', False):
                # 如果价格回到开仓价格，平掉第一次加仓的仓位
                if current_price <= position['entry_price']:
                    return current_price, "平掉第一次加仓"
        
        return None, ""

    def _calculate_pnl_ratio(self, position, current_price):
        """计算盈亏比例"""
        # 从 FUTURES_CONFIG 获取配置
        future_config = FUTURES_CONFIG.get(position['pair_name'])
        point_value = future_config['point_value']
        
        if position['direction'] == '多头':
            pnl = (current_price - position['entry_price']) * position['lots'] * point_value
        else:
            pnl = (position['entry_price'] - current_price) * position['lots'] * point_value
            
        return pnl / self.capital

    def _close_position(self, position, exit_price, exit_date, reason=""):
        """平仓"""
        future_config = FUTURES_CONFIG.get(position['pair_name'])
        point_value = future_config['point_value']
        
        if position['direction'] == '多头':
            pnl = (exit_price - position['entry_price']) * position['lots'] * point_value
        else:
            pnl = (position['entry_price'] - exit_price) * position['lots'] * point_value
        
        # 释放保证金
        self.available_capital += abs(position['margin'])  # 确保为正数
        self.capital += pnl
        
        return {
            '记录类型': '平仓',
            '开仓日期': position['entry_date'],
            '平仓日期': exit_date,
            '方向': position['direction'],
            '开仓价格': position['entry_price'],
            '平仓价格': exit_price,
            '交易手数': position['lots'],
            '盈亏': pnl,
            '平仓原因': reason,
            '收益率': pnl / self.initial_capital,
            '保证金': abs(position['margin']),  # 确保为正数
            '当前权益': self.capital
        }
    
    def _calculate_realized_pnl(self, position, exit_price):
        """计算已实现盈亏"""
        if position['direction'] == '多头':
            return position['size'] * (exit_price - position['entry_price']) / position['entry_price']
        else:
            return position['size'] * (position['entry_price'] - exit_price) / position['entry_price']
    
    def _calculate_unrealized_pnl(self, position, current_price):
        """计算未实现盈亏"""
        # 获取期货配置
        future_config = FUTURES_CONFIG.get(position['pair_name'])
        point_value = future_config['point_value']
        
        if position['direction'] == '多头':
            return (current_price - position['entry_price']) * position['lots'] * point_value
        else:
            return (position['entry_price'] - current_price) * position['lots'] * point_value
    
    def _generate_backtest_report(self, trades, equity_curve, pair_name):
        """生成回测报告"""
        if not trades:
            return {
                '品种名称': pair_name,
                '总交易次数': 0,
                '盈利次数': 0,
                '胜率': 0,
                '总盈亏': 0,
                '最大回撤': 0,
                '夏普比率': 0,
                '权益曲线': pd.DataFrame(equity_curve),
                '交易记录': pd.DataFrame()  # 空的交易记录
            }
        
        df_trades = pd.DataFrame(trades)
        df_equity = pd.DataFrame(equity_curve)
        
        # 计算关键指标
        total_trades = len(trades)
        profitable_trades = len(df_trades[df_trades['盈亏'] > 0])
        win_rate = profitable_trades / total_trades
        total_pnl = df_trades['盈亏'].sum()
        
        # 计算最大回撤
        df_equity['回撤'] = df_equity['equity'].cummax() - df_equity['equity']
        max_drawdown = df_equity['回撤'].max()
        
        # 计算夏普比率
        returns = df_equity['equity'].pct_change().dropna()
        sharpe_ratio = np.sqrt(252) * returns.mean() / returns.std()
        
        return {
            '品种名称': pair_name,
            '总交易次数': total_trades,
            '盈利次数': profitable_trades,
            '胜率': win_rate,
            '总盈亏': total_pnl,
            '最大回撤': max_drawdown,
            '夏普比率': sharpe_ratio,
            '权益曲线': df_equity,
            '交易记录': df_trades
        }
    
    def plot_results(self, results):
        """绘制回测结果"""
        try:
            fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(15, 10))
            
            # 绘制权益曲线
            equity_curve = results['权益曲线']
            ax1.plot(equity_curve['date'], equity_curve['equity'])
            
            # 计算收益率
            final_equity = equity_curve['equity'].iloc[-1]
            total_return = (final_equity - self.initial_capital) / self.initial_capital * 100
            
            ax1.set_title(f'权益曲线 - {results["品种名称"]}\n'
                         f'初始资金: {self.initial_capital:,.0f}, 最终资金: {final_equity:,.0f}\n'
                         f'总收益率: {total_return:.2f}%', 
                         fontproperties=self.font_prop)
            ax1.set_xlabel('日期', fontproperties=self.font_prop)
            ax1.set_ylabel('权益', fontproperties=self.font_prop)
            ax1.grid(True)
            
            # 设置x轴日期格式
            ax1.tick_params(axis='x', rotation=45)
            
            # 绘制交易分布
            if len(results['交易记录']) > 0:
                results['交易记录']['盈亏'].hist(ax=ax2, bins=50)
                ax2.set_title('盈亏分布', fontproperties=self.font_prop)
                ax2.set_xlabel('盈亏', fontproperties=self.font_prop)
                ax2.set_ylabel('频率', fontproperties=self.font_prop)
                ax2.grid(True)
            
            # 添加cyberpunk风格
            mplcyberpunk.make_lines_glow()
            plt.tight_layout()
            return fig
        except Exception as e:
            print(f"绘图出错: {str(e)}")
            return None 

    def _record_position_status(self, position, current_price, date):
        """记录持仓状态"""
        return {
            '记录类型': '持仓',
            '日期': date,
            '方向': position['direction'],
            '持仓价格': position['entry_price'],
            '当前价格': current_price,
            '持仓手数': position['lots'],
            '止损价': "——",
            '止盈价1': "——",
            '止盈价2': "——",
            '保证金': abs(position['margin']),  # 确保为正数
            '浮动盈亏': self._calculate_unrealized_pnl(position, current_price)
        }

    def _record_open_position(self, position):
        """记录开仓信息"""
        return {
            '记录类型': '开仓',
            '日期': position['entry_date'],
            '方向': position['direction'],
            '开仓价格': position['entry_price'],
            '交易手数': position['lots'],
            '网格大小': position['grid_size'],  # 改用网格大小
            '保证金': abs(position['margin']),  # 确保为正数
            '初始权益': self.initial_capital,
            '加仓状态': '初始开仓'  # 添加加仓状态说明
        }

    def calculate_current_equity(self, position, current_price):
        """计算当前权益"""
        equity = self.capital
        if position:
            unrealized_pnl = self._calculate_unrealized_pnl(position, current_price)
            equity += unrealized_pnl
        return equity 

    def _handle_position_adjustment(self, position, current_price, action, date):
        """处理仓位调整（加仓或减仓）"""
        trade_record = None  # 初始化交易记录
        
        if "加仓" in action:
            # 计算新的开仓手数
            new_lots = self.calculate_position_size(
                self.calculate_current_equity(position, current_price),
                current_price,
                position['pair_name']
            )
            
            if new_lots > 0:
                # 更新持仓信息
                if "第一次" in action:
                    position['first_add'] = True
                    position['first_add_price'] = current_price
                    position['first_add_lots'] = new_lots
                else:  # 第二次加仓
                    position['second_add'] = True
                    position['second_add_price'] = current_price
                    position['second_add_lots'] = new_lots
                
                # 更新总持仓手数
                position['lots'] += new_lots
                
                # 更新保证金
                future_config = FUTURES_CONFIG.get(position['pair_name'])
                margin_ratio = future_config['margin_ratio']
                point_value = future_config['point_value']
                new_margin = current_price * point_value * new_lots * margin_ratio
                position['margin'] += new_margin
                self.available_capital -= new_margin
                
                # 记录加仓交易
                trade_record = {
                    '记录类型': '加仓',
                    '日期': date,
                    '方向': position['direction'],
                    '加仓价格': current_price,
                    '加仓手数': new_lots,
                    '加仓类型': "第一次加仓" if "第一次" in action else "第二次加仓",
                    '保证金': new_margin,
                    '当前权益': self.calculate_current_equity(position, current_price)
                }
        
        elif "平掉" in action:
            # 计算需要平仓的手数和对应的开仓价格
            if "第二次" in action:
                lots_to_close = position['second_add_lots']
                entry_price = position['second_add_price']
                position['second_add'] = False
                position['lots'] -= lots_to_close
            else:  # 平掉第一次加仓
                lots_to_close = position['first_add_lots']
                entry_price = position['first_add_price']
                position['first_add'] = False
                position['lots'] -= lots_to_close
            
            # 计算平仓盈亏
            future_config = FUTURES_CONFIG.get(position['pair_name'])
            point_value = future_config['point_value']
            if position['direction'] == '多头':
                pnl = (current_price - entry_price) * lots_to_close * point_value
            else:
                pnl = (entry_price - current_price) * lots_to_close * point_value
            
            # 更新资金
            self.capital += pnl
            margin_released = position['margin'] * (lots_to_close / position['lots'])
            self.available_capital += margin_released
            position['margin'] *= (1 - lots_to_close / position['lots'])
            
            # 记录平仓交易
            trade_record = {
                '记录类型': '部分平仓',
                '日期': date,
                '方向': position['direction'],
                '开仓价格': entry_price,
                '平仓价格': current_price,
                '平仓手数': lots_to_close,
                '平仓类型': "平掉第二次加仓" if "第二次" in action else "平掉第一次加仓",
                '盈亏': pnl,
                '释放保证金': margin_released,
                '当前权益': self.calculate_current_equity(position, current_price)
            }
        
        return position, trade_record 